import torch.nn as nn
from torch.nn import functional as F


F_MARGIN = 0.2


class ZFNet(nn.Module):
    def __init__(self):
        super(ZFNet, self).__init__()
        # n*3*112*112
        self.c1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.ELU(inplace=True))
        # n*64*56*56
        self.c2 = nn.Sequential(
            nn.Conv2d(64, 256, kernel_size=5, stride=2, padding=2),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ELU(inplace=True)
        )
        # n*256*14*14
        self.c3_5 = nn.Sequential(
            nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
            nn.ELU(inplace=True),
            nn.Conv2d(256, 64, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ELU(inplace=True))
        # n*64*7*7 -》n*3136
        self.fc6_8 = nn.Sequential(
            nn.Linear(64 * 7 * 7, 2048),
            nn.Softsign(),
            nn.Linear(2048, 512),
            nn.Softsign(),
            nn.Linear(512, 128))
        # n*128
        self.f_thres = 0.5

    def forward(self, x):
        x = self.c1(x)
        x = self.c2(x)
        x = self.c3_5(x)
        x = x.view(-1, 64 * 7 * 7)
        x = self.fc6_8(x)
        x = F.normalize(x, p=2, dim=1)
        return x

    def MyLoss(self, x):
        x = x.view(x.size(0) // 3, 3, -1)
        x = x.permute((1, 0, 2)).contiguous()
        dis_ap = (x[0] - x[1]).pow(2).sum(1)
        dis_an = (x[0] - x[2]).pow(2).sum(1)
        self.f_thres = 0.9 * self.f_thres + 0.1 * (dis_an.mean() + dis_ap.mean()) / 2.0
        loss = F.relu(dis_ap - dis_an + F_MARGIN).sum()
        return loss

    def MyCor(self, nn_outputs, nn_flags):
        nn_outdis = nn_outputs.view(nn_outputs.size(0) // 2, 2, -1)
        nn_outdis = nn_outdis.permute((1, 0, 2)).contiguous()
        nn_outdis = (nn_outdis[0] - nn_outdis[1]).pow(2).sum(1)
        nn_outnear = (nn_outdis < self.f_thres)
        return (nn_outnear == nn_flags)
